#include <butil/string_printf.h>

#include "index/inverted_index/inverted_disk_indexer/inverted_disk_indexer.h"
#include "util/file_wrapper/file_wrapper.h"

namespace hawking {
namespace indexlib {

DEFINE_int32(term_expand_factor, 10, "max term count could in realtime");

Status InvertedDiskIndexer::Load(
    const std::string& index_dir, bool need_increase,
    FSFileType data_filetype, size_t doc_count) {
    need_increase_ = need_increase;
    doc_count_ = doc_count;

    std::string offset_file = butil::string_printf("%s/offset_file", index_dir.data());
    offset_file_reader_ = std::make_unique<MmapFileReader>(offset_file, need_increase);
    Status ret = offset_file_reader_->Init();
    if (Status::OK != ret) {
        LOG(ERROR) << "offset_file_reader Init " << offset_file
            << " Fail with " << static_cast<int32_t>(ret);
        return ret;
    }

    std::string data_file = butil::string_printf("%s/data_file", index_dir.data());
    if (data_filetype == FSFileType::FSFT_MMAP) {
        data_file_reader_ = std::make_unique<MmapFileReader>(data_file);
        Status ret = data_file_reader_->Init();
        if (ret != Status::OK) {
            LOG(ERROR) << "mmap file reader Init fail with "
                << static_cast<int32_t>(ret);
            return ret;
        }
    } else {
        data_file_reader_ = std::make_unique<BlockFileReader>(data_file);
        Status ret = data_file_reader_->Init();
        if (ret != Status::OK) {
            LOG(ERROR) << "blockcache file reader Init fail with "
                << static_cast<int32_t>(ret);
            return ret;
        }
    }
    ret = offset_file_reader_->Init();
    if (Status::OK != ret) {
        LOG(ERROR) << "data_file_reader Init " << data_file
            << " Fail with " << static_cast<int32_t>(ret);
        return ret;
    }

    fix_term_count_ = offset_file_reader_->Length() / offset_file_item_bytes_;
    offset_file_reader_->Read((uint8_t*)&min_term_key_, sizeof(min_term_key_), 0);
    offset_file_reader_->Read(
        (uint8_t*)&max_term_key_,
        sizeof(max_term_key_),
        offset_file_reader_->Length() - offset_file_item_bytes_);

    if (need_increase) {
        terms_updator_.resize(fix_term_count_ * FLAGS_term_expand_factor);
    }
    return Status::OK;
}

Status InvertedDiskIndexer::RemoveDocAtDocUpdator_(uint64_t old_term, DocID doc_id) {
    InvertedTermUpdator* docs_updator = nullptr;
    FindDocUpdator_(old_term, &docs_updator);
    if (!docs_updator) {
        return Status::NotFound;
    }

    docs_updator->Clear(doc_id);
    return Status::OK;
}

Status InvertedDiskIndexer::AddDocAtDocUpdator_(uint64_t new_term, DocID doc_id) {
    InvertedTermUpdator* docs_updator = nullptr;
    FindDocUpdator_(new_term, &docs_updator);
    if (!docs_updator) {
        uint64_t offset_in_doc_updator;
        InvertedTermUpdator* new_term_updator =
            FetchDocUpdator_(new_term, &offset_in_doc_updator);
        if (!new_term_updator) {
            return Status::NoMem;
        }

        new_term_updator->Set(doc_id);
        new_term_updator->SetValid();
    }

    docs_updator->Set(doc_id);
    return Status::OK;
}

Status InvertedDiskIndexer::RemoveDocAtOldTerm_(
    uint64_t old_term, DocID rem_doc_id, SegmentContext* segment_ctx) {
    uint64_t offset_in_datafile;
    uint32_t doclist_length;
    // old_term maybe realtime-increase term after load
    if (!FindOffsetInOffsetFile_(old_term, &offset_in_datafile, &doclist_length)) {
        return RemoveDocAtDocUpdator_(old_term, rem_doc_id);
    }

    // old_term is from fix-datafile when load
    if (offset_in_datafile >= data_file_reader_->Length()) {
        // not first updated
        InvertedTermUpdator* docs_updator = nullptr;
        GetDocUpdator_(offset_in_datafile - data_file_reader_->Length(), &docs_updator);
        if (unlikely(!docs_updator)) {
            return Status::NotFound;
        }

        docs_updator->Clear(rem_doc_id);
        return Status::OK;
    } else {
        // first updated. load current doclist
        if (unlikely(Status::OK != Lookup_(offset_in_datafile, INT_MAX, segment_ctx, false))) {
            return Status::IOError;
        }

        // fetch a doc updator, set doclist without remove doc_id
        uint64_t offset_in_doc_updator = 0;
        InvertedTermUpdator* docs_updator = FetchDocUpdator_(old_term, &offset_in_doc_updator);
        if (!docs_updator) {
            return Status::NoMem;
        }

        for (auto doc_id : segment_ctx->DocIds()) {
            if (doc_id != rem_doc_id) {
                docs_updator->Set(doc_id);
            }
        }
        docs_updator->SetValid();

        // update offset_in_datafile to offsetfile with `datafile length + doc_updator_idx`
        size_t new_offset = offset_in_doc_updator + data_file_reader_->Length();
        size_t offset_in_offsetfile = (old_term - min_term_key_) * offset_file_item_bytes_;
        return offset_file_reader_->Write(
            (const uint8_t*)&new_offset, sizeof(new_offset), offset_in_offsetfile).Code();
    }
}

Status InvertedDiskIndexer::AddDocAtNewTerm_(
    uint64_t new_term, DocID add_doc_id, SegmentContext* segment_ctx) {
    uint64_t offset_in_datafile;
    uint32_t doclist_length;
    // new_term is realtime-increase term after load
    if (!FindOffsetInOffsetFile_(new_term, &offset_in_datafile, &doclist_length)) {
        return AddDocAtDocUpdator_(new_term, add_doc_id);
    }

    if (offset_in_datafile >= data_file_reader_->Length()) {
        // new_term is in doc updator
        InvertedTermUpdator* docs_updator = nullptr;
        GetDocUpdator_(offset_in_datafile - data_file_reader_->Length(), &docs_updator);
        docs_updator->Set(add_doc_id);
    } else {
        // new_term is in datafile
        // first updated. load current doclist
        if (unlikely(Status::OK != Lookup_(offset_in_datafile, INT_MAX, segment_ctx, false))) {
            return Status::IOError;
        }

        // fetch a doc updator, set doclist without remove doc_id
        uint64_t offset_in_doc_updator = 0;
        InvertedTermUpdator* docs_updator = FetchDocUpdator_(new_term, &offset_in_doc_updator);
        if (!docs_updator) {
            return Status::NoMem;
        }

        segment_ctx->MutableDocIds()->emplace_back(add_doc_id);
        for (auto doc_id : segment_ctx->DocIds()) {
            docs_updator->Set(doc_id);
        }

        // update offset_in_datafile to offsetfile with `datafile length + doc_updator_idx`
        size_t new_offset = offset_in_doc_updator + data_file_reader_->Length();
        size_t offset_in_offsetfile = (new_term - min_term_key_) * offset_file_item_bytes_;
        return offset_file_reader_->Write(
            (const uint8_t*)&new_offset, sizeof(new_offset), offset_in_offsetfile).Code();
    }
}

Status InvertedDiskIndexer::UpdateDocument(
    uint64_t old_term, uint64_t new_term, DocId doc_id) {
    SegmentContext segment_ctx;
    ret = RemoveDocAtOldTerm_(old_term, doc_id, &segment_ctx);
    if (ret != Status::OK) {
        LOG(ERROR) << "remove doc_id " << doc_id << " from term " << old_term << " fail";
        return ret;
    }

    segment_ctx.Clear();
    return AddDocAtNewTerm_(new_term, doc_id, &segment_ctx);
}

void InvertedDiskIndexer::GetDocUpdator_(
    uint64_t offset_in_doc_updator, InvertedTermUpdator** docs_updator) {
    if (terms_updator_[offset_in_doc_updator].IsValid()) {
        *docs_updator = &terms_updator_[offset_in_doc_updator];
    }
}

void InvertedDiskIndexer::FindDocUpdator_(
    uint64_t term_key, InvertedTermUpdator** docs_updator) {
    for (size_t idx = 0; idx < terms_updator_.size(); ++idx) {
        if (!terms_updator_[idx].IsValid()) {
            return;
        }

        if (terms_updator_[idx].TermKey() == term_key) {
            *docs_updator = terms_updator_[idx].get();
            return;
        }
    }
}

InvertedTermUpdator* InvertedDiskIndexer::FetchDocUpdator_(
    uint64_t term, uint64_t* offset_in_doc_updator) {
    if (dynamic_term_count_ == terms_updator_.size()) {
        return nullptr;
    }

    *offset_in_doc_updator = dynamic_term_count_++;
    terms_updator_[*offset_in_doc_updator].Init(term, DocCount());
    return &terms_updator_[*offset_in_doc_updator];
}

Status InvertedDiskIndexer::LookupInTermUpdator_(
    uint32_t topk, uint64_t term_key, SegmentContext8 segment_ctx) {
    InvertedTermUpdator* docs_updator = nullptr;
    FindDocUpdator_(term_key, &docs_updator);
    if (!docs_updator) {
        return Status::NotFound;
    }

    docs_updator->TraverseTopk(segment_ctx->MutableDocIds(), topk);
    for (size_t idx = 0; idx < segment_ctx->DocIds().size(); ++idx) {
        segment_ctx->MutableDocContext()->emplace_back();
    }

    return Status::OK;
}

Status InvertedDiskIndexer::LookupInTermUpdator_(
    uint64_t offset_in_datafile, uint32_t topk, SegmentContext* segment_ctx) {
    InvertedTermUpdator* docs_updator = nullptr;
    GetDocUpdator_(offset_in_datafile - data_file_reader_->Length(), &docs_updator);
    if (unlikely(!docs_updator)) {
        return Status::OutOfRange;
    }

    docs_updator->TraverseTopk(segment_ctx->MutableDocIds(), topk);
    for (size_t idx = 0; idx < segment_ctx->DocIds().size(); ++idx) {
        segment_ctx->MutableDocContext()->emplace_back();
    }

    return Status::OK;
}

Status InvertedDiskIndexer::Lookup_(
    uint64_t offset_in_datafile,
    uint32_t topk,
    SegmentContext* segment_ctx,
    bool need_doc_context) {
    if (!topk) {
        return Status::OK;
    }

    segment_ctx->Reserve(topk);
    if (offset_in_datafile >= data_file_reader_->Length()) {
        return LookupInTermUpdator_(offset_in_datafile, topk, segment_ctx);
    }

    Status ret = data_file_reader_->Read(
        (uint8_t*)segment_ctx->MutableDocIds(),
        topk * sizeof(DocId),
        offset_in_datafile).Code();
    if (ret != Status::OK) {
        segment_ctx->Clear();
    }

    if (need_doc_context) {
        segment_ctx->MutableDocContext()->resize(topk);
        segment_ctx->SetRelateDocCount(topk);
    }
    return ret;
}

bool InvertedDiskIndexer::FindOffsetInOffsetFile_(
    const uint64_t term_key, uint64_t* offset_in_datafile, uint32_t* doclist_length) {
    if (term_key < min_term_key_ || term_key > max_term_key_) {
        return false;
    }

    size_t low = 0;
    size_t high = offset_file_reader_->Length() / offset_file_item_bytes_ - 1;

    while (low <= high) {
        size_t mid = low + ((high - low) >> 1);
        uint64_t mid_term_key;
        offset_file_reader_->Read(
            (uint8_t*)&mid_term_key, sizeof(mid_term_key), mid * offset_file_item_bytes_);

        if (mid_term_key == term_key) {
            offset_file_reader_->Read(
                (uint8_t*)offset_in_datafile,
                sizeof(*offset_in_datafile),
                mid * offset_file_item_bytes_ + term_key_length_);
            offset_file_reader_->Read(
                (uint8_t*)doclist_length,
                sizeof(*doclist_length),
                mid * offset_file_item_bytes_ + term_key_length_ + offset_in_datafile_length_);
            return true;
        } else if (mid_term_key < term_key) {
            low = mid + 1;
        } else {
            high = mid - 1;
        }
    }

    return false;
}

Status InvertedDiskIndexer::Lookup(
    uint64_t term_key, uint32_t topk, SegmentContext* segment_ctx) {
    uint64_t offset_in_datafile;
    uint32_t doclist_length;
    if (!FindOffsetInOffsetFile_(term_key, &offset_in_datafile, &doclist_length)) {
        if (!need_increase_) {
            return Status::NotFound;
        }

        return LookupInTermUpdator_(term_key, topk, segment_ctx);
    }
    
    return Lookup_(offset_in_datafile, std::min(topk, doclist_length), segment_ctx);
}

}
}